package com.hhj.sharding.sphere.demo.algorithm;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;

import java.util.Collection;
import java.util.HashSet;
import java.util.Set;

/**
 * 自定义复合分库策略算法
 *
 * @author HHJ
 * @since 2024-05-23
 */
@Slf4j
public class CustomDatabaseComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm<Long> {
    /**
     * @param dataSourceNames 库名集合
     * @param shardingValue   分片键信息
     * @return 匹配的数据源集合
     */
    @Override
    public Collection<String> doSharding(Collection<String> dataSourceNames, ComplexKeysShardingValue<Long> shardingValue) {
        Collection<Long> orderIds = shardingValue.getColumnNameAndShardingValuesMap().get("user_id");
        Collection<Long> otherIds = shardingValue.getColumnNameAndShardingValuesMap().get("other_id");

        if (CollectionUtils.isNotEmpty(orderIds)) {
            return getDatabaseNames(dataSourceNames, orderIds);
        }

        if (CollectionUtils.isNotEmpty(otherIds)) {
            return getDatabaseNames(dataSourceNames, otherIds);
        }

        return null;
    }

    /**
     * 获取数据源名
     *
     * @param columnValues 分片键的值
     * @return 数据源集合
     */
    private Set<String> getDatabaseNames(Collection<String> dataSourceNames, Collection<Long> columnValues) {
        Set<String> set = new HashSet<>();
        for (Long columnValue : columnValues) {
            // 一般不会使用这些ID直接去分库，不然即使你同一个用户，也将将随机将数据分配到各个数据库，这合理吗？
            // 明显不合理，那到时候关联查询怎么办？虽然sharding sphere 是支持查询的，但是这是笛卡尔积的查询，是非常消耗性能的
            // 所以为设计ID的时候，一般就将ID的最后几位换成用户ID，这样，这里直接获取到后面几位的用户ID
            // 这样就能固定这个ID了，到时候所有的数据将会全部路由到同一个数据库
            // 我们这里是直接分库的，仅仅只是演示
            String databaseSuffix = ShardingAlgorithmHelper.getDatabaseSuffix(columnValue);
            log.info("database suffix {}", databaseSuffix);
            for (String dataSourceName : dataSourceNames) {
                if (dataSourceName.endsWith(databaseSuffix)) {
                    set.add(dataSourceName);
                }
            }
        }
        return set;
    }


}
